# from paddlex import create_model
# import numpy as np
# import sys
# import os
# # 获取到 config.py 所在的目录（相对于 get_table_cell.py 的路径）
# current_dir = os.path.dirname(os.path.abspath(__file__))
# config_dir = os.path.abspath(os.path.join(current_dir, '../../..'))  # 返回到 src 目录

# # 将 config.py 所在的路径添加到 sys.path
# sys.path.append(config_dir)
# from config import config
# LABEL_NAME = config.label_name
# THRESHOLD = config.threshold


# def get_table_label(img) -> str:
#     """
#     根据图片路径预测表格的类型，返回表格标签（'wired_table' 或 'wireless_table'）。
#     :param img: Python Var/str/list,
#         Python变量，如numpy.ndarray表示的图像数据
#         文件路径，如图像文件的本地路径：/root/data/img.jpg
#         URL链接，如图像文件的网络URL：示例
#         本地目录，该目录下需包含待预测数据文件，如本地路径：/root/data/
#         列表，列表元素需为上述类型数据，如[numpy.ndarray, numpy.ndarray]，[\"/root/data/img1.jpg\", \"/root/data/img2.jpg\"]，[\"/root/data1\", \"/root/data2\"]
#     """
#     label_model = create_model(model_name="PP-LCNet_x1_0_table_cls")
#     label_output = label_model.predict(img, batch_size=1)
#     label_output = list(label_output)[0]
    
#     # 获取分数最高的标签
#     max_score_index = np.argmax(label_output['scores'])
#     label_name = label_output['label_names'][max_score_index]
    
#     return label_name



# def get_cell_coordinates(img, label_name: str = LABEL_NAME, threshold: float = THRESHOLD) -> list:
#     """
#     根据表格标签获取表格单元格的坐标列表。
    
#     :param img: Python Var/str/list,
#         Python变量，如numpy.ndarray表示的图像数据
#         文件路径，如图像文件的本地路径：/root/data/img.jpg
#         URL链接，如图像文件的网络URL：示例
#         本地目录，该目录下需包含待预测数据文件，如本地路径：/root/data/
#         列表，列表元素需为上述类型数据，如[numpy.ndarray, numpy.ndarray]，[\"/root/data/img1.jpg\", \"/root/data/img2.jpg\"]，[\"/root/data1\", \"/root/data2\"]
#     :param label_name: 表格标签 ('wired_table' 或 'wireless_table')
#     :param threshold: 用于预测时过滤的置信度阈值，默认为0.3
#     :return: 包含单元格坐标的列表
#     """
#     if label_name not in ['wired_table', 'wireless_table','auto']:
#         raise ValueError(f"Invalid label_name '{label_name}'. Expected 'wired_table' 、'auto' or 'wireless_table'.")
    
#     if label_name == 'auto':
#         label_name = get_table_label(img)
#     if label_name == 'wired_table':
#         model = create_model(model_name="RT-DETR-L_wired_table_cell_det")
#     else:
#         model = create_model(model_name="RT-DETR-L_wireless_table_cell_det")
    
#     # 获取单元格坐标
#     cell_coords = model.predict(img, threshold=threshold, batch_size=1)
#     cell_boxes = list(cell_coords)[0]['boxes']
    
#     # 提取坐标
#     cell_coordinates = [boxe['coordinate'] for boxe in cell_boxes]
    
#     return cell_coordinates

from paddlex import create_model
import numpy as np
import sys
import os
# 获取到 config.py 所在的目录（相对于 get_table_cell.py 的路径）
current_dir = os.path.dirname(os.path.abspath(__file__))
config_dir = os.path.abspath(os.path.join(current_dir, '../../..'))  # 返回到 src 目录

# 将 config.py 所在的路径添加到 sys.path
sys.path.append(config_dir)
from config import config
LABEL_NAME = config.label_name
THRESHOLD = config.threshold

class ModelRunner:
    def __init__(self, label_model_name="PP-LCNet_x1_0_table_cls", cell_model_wired="RT-DETR-L_wired_table_cell_det", cell_model_wireless="RT-DETR-L_wireless_table_cell_det"):
        """
        初始化模型加载，仅加载一次
        :param label_model_name: 标签预测模型名称
        :param cell_model_wired: 有线表格单元格检测模型名称
        :param cell_model_wireless: 无线表格单元格检测模型名称
        """
        self.label_model = create_model(model_name=label_model_name)  # 标签预测模型
        self.cell_model_wired = create_model(model_name=cell_model_wired)  # 有线表格单元格检测模型
        self.cell_model_wireless = create_model(model_name=cell_model_wireless)  # 无线表格单元格检测模型

    def get_table_label(self, img) -> str:
        """
        根据图片路径预测表格的类型，返回表格标签（'wired_table' 或 'wireless_table'）。
        :param img: 输入图像
        :return: 表格标签
        """
        label_output = self.label_model.predict(img, batch_size=1)
        label_output = list(label_output)[0]
        
        # 获取分数最高的标签
        max_score_index = np.argmax(label_output['scores'])
        label_name = label_output['label_names'][max_score_index]
        
        return label_name

    def get_cell_coordinates(self, img, label_name: str=LABEL_NAME, threshold: float = THRESHOLD) -> list:
        """
        根据表格标签获取表格单元格的坐标列表。
        :param img: 输入图像
        :param label_name: 表格标签 ('wired_table' 或 'wireless_table')
        :param threshold: 用于预测时过滤的置信度阈值，默认为0.3
        :return: 包含单元格坐标的列表
        """
        if label_name not in ['wired_table', 'wireless_table', 'auto']:
            raise ValueError(f"Invalid label_name '{label_name}'. Expected 'wired_table' or 'wireless_table' or 'auto'.")

        if label_name == 'auto':
            label_name = self.get_table_label(img)
        
        # 根据标签选择模型
        if label_name == 'wired_table':
            model = self.cell_model_wired
        else:
            model = self.cell_model_wireless
        
        # 获取单元格坐标
        cell_coords = model.predict(img, threshold=threshold, batch_size=1)
        cell_boxes = list(cell_coords)[0]['boxes']
        
        # 提取坐标
        cell_coordinates = [boxe['coordinate'] for boxe in cell_boxes]
        
        return cell_coordinates
